#!/bin/env pypy

"""
Tools for the manipulation of VCF files.

Usage:
  vcf filter <vcf_file> [options]
  vcf summarize <vcf_file> [options]

Options:
  -h --help                 Show this screen.
  --min-qual=N              Minimum genotype quality [default: 100].
  --min-alt-control-qual=N  Minimum genotype quality for alternate alleles in
						    controls [default: 60].
  --test=REGEXP             Regexp for test sample names.
  --control=REGEXP          Regexp for control sample names.
  --abs-negative=REGEXP     Regexp for absolute negative sample names.
  --nonsynonymous           Only keep non-synonymous coding variants.
  --min-alt=N               Minimum number of alternate alleles [default: 0].
  --baseq-alpha=FLOAT       Base quality bias p-value threshold [default: 0].

  --group-names=ID,...      Sample group identifiers.
  --group-regex=REGEX,...   Regular expressions that will be matched against
                            samples to identify samples in a particular group.
  --rank=METHOD             Method used for ranking variants. Available options:
							frequency, gene_frequency [default: frequency].
"""

from __future__ import print_function
import sys, subprocess, docopt, re, os
from collections import defaultdict
from pypette import zopen, shell, argsort
from pypette import info, error










##############
# VCF FILTER #
##############

def vcf_filter(vcf_path, min_qual, min_alt_control_qual,
	test_regex, control_regex, baseq_p_threshold,
	only_nonsynonymous, min_alternate):
	
	if min_alt_control_qual > min_qual:
		error('--min-alt-control-qual must be less than --min-qual.')
	
	vcf_file = zopen(vcf_path)
	for line in vcf_file:
		if not line.startswith('##'): break
	
	headers = line[:-1].split('\t')
	sample_col = headers.index('FORMAT') + 1
	
	if only_nonsynonymous:
		col_synonymous = headers.index('SYNONYMOUS')
	
	samples = headers[sample_col:]
	samples = [re.sub(r'(.*/)?(.*).bam', r'\2', s) for s in samples]
	
	print('\t'.join(headers[:sample_col] + samples))
	
	test = [True] * len(samples)
	if test_regex:
		test = [re.search(test_regex, s, re.I) != None for s in samples]
		
	control = [False] * len(samples)
	if control_regex:
		control = [re.search(control_regex, s, re.I) != None for s in samples]
	have_controls = any(control)
		
	gt_options = ['0/0', '0/1', '1/1']
	
	pv4_regex = re.compile('PV4=(\S+),(\S+),(\S+),(\S+)')
	
	for line in vcf_file:
		cols = line[:-1].split('\t')
		gt_cols = cols[sample_col:]
		
		# Discard variants that exhibit base quality bias.
		# This check is disabled for indels, because we have observed indels
		# that validate fine but have a huge baseQ bias according to samtools.
		if baseq_p_threshold > 0:
			m = pv4_regex.search(cols[7])
			baseq_p = float(m.group(2)) if m else 1
			if baseq_p < baseq_p_threshold and cols[7].find('INDEL') == -1:
				continue
		
		# Discard all variants except non-synonymous ones if requested.
		if only_nonsynonymous:
			if re.search('^(synonymous|unknown|-)', cols[col_synonymous]):
				continue
			
		genotypes = [gt_options.index(g[0:3]) for g in gt_cols]
		quals = [g.split(':')[1].split(',')[0:2] for g in gt_cols]
		quals = [max(int(q[0]), int(q[1])) for q in quals]
		
		# Set genotype as unknown (-1) for non-control samples and non-mutated
		# samples if quality is below the minimum threshold.
		genotypes = [-1 if (not (control[k] and g > 0)) and
			quals[k] < min_qual else g for k, g in enumerate(genotypes)]
		
		# Set genotype as unknown (-1) for mutated control samples if
		# quality is below the minimum threshold for this type of sample.
		genotypes = [-1 if control[k] and g > 0 and
			quals[k] < min_alt_control_qual else g
			for k, g in enumerate(genotypes)]
		
		# Discard variant if any absolute negative samples show alt. alleles.
		if any(g > 0 and control[k] for k, g in enumerate(genotypes)):
			continue
		
		# Discard variant if there are not sufficiently many alternate alleles
		# among the test cohort.
		if sum(g > 0 and test[k]
			for k, g in enumerate(genotypes)) < min_alternate:
			continue
		
		# Discard variant if all control samples could not be reliably
		# genotyped, or if all test samples are negative or unknown.
		if all(g < 1 for k, g in enumerate(genotypes) if test[k]): continue
		
		if have_controls:
			if all(g == -1 for k, g in enumerate(genotypes) if control[k]):
				continue
			
		sys.stdout.write(line)
		







#################
# VCF SUMMARIZE #
#################

def vcf_summarize(vcf_path, min_qual=100, group_names=None, group_regex=None,
	rank_method='frequency'):
	
	vcf_file = open(vcf_path)
	for line in vcf_file:
		if not line.startswith('##'): break
	
	headers = line[:-1].split('\t')
	sample_col = headers.index('FORMAT') + 1
	col_cosmic = headers.index('COSMIC')
	col_1000g = headers.index('1000GENOMES')
	col_nearby_genes = headers.index('NEARBY_GENES')
	
	samples = headers[sample_col:]
	samples = [re.sub(r'(.*/)?(.*).bam', r'\2', s) for s in samples]
	
	sample_order = range(len(samples))
	if group_names and group_regex:
		# User wants samples grouped, so we construct a new sample order.
		group_names = group_names.split(',')
		group_regex = group_regex.split(',')
		group_samples = []
		
		if len(group_names) != len(group_regex):
			error('Number of group names and regular expressions must match.')
		
		sample_order = []
		for regex in group_regex:
			orig_indices = [k for k, s in enumerate(samples) if
				re.search(regex, s, re.I)]
			group_samples.append(range(len(sample_order),
				len(sample_order) + len(orig_indices)))
			sample_order += orig_indices
	else:
		group_names = ['TOTAL']
		group_samples = [range(len(samples))]
	
	if len(sample_order) < len(samples):
		info('WARNING: Some samples were not placed in any sample group.')
	
	samples = [samples[i] for i in sample_order]
	
	gt_options = ['0/0', '0/1', '1/1']
	
	variant_lines = []
	variant_scores = []
	gene_variants = defaultdict(list)
	
	for line in vcf_file:
		cols = line[:-1].split('\t')
		info_cols = cols[:sample_col]
		gt_cols = cols[sample_col:]
		gt_cols = [gt_cols[i] for i in sample_order]
		
		genotypes = [gt_options.index(g[0:3]) for g in gt_cols]
		quals = [g.split(':')[1].split(',')[0:2] for g in gt_cols]
		quals = [max(int(q[0]), int(q[1])) for q in quals]
		
		# Set genotype as unknown (-1) if quality is below a threshold.
		genotypes = [-1 if quals[k] < min_qual else g
			for k, g in enumerate(genotypes)]
		
		# Clear out unknown genotypes.
		gt_cols = [' ' if g == -1 else '%s (%d)' % (gt_options[g], quals[k])
			for k, g in enumerate(genotypes)]
		
		# Calculate total number of mutated samples in each sample group.
		# FIXME: Rewrite using numpy fancy indexing.
		group_homozygous = [
			sum([g >= 2 for g in [genotypes[s] for s in _samples]])
			for _samples in group_samples]
		group_heterozygous = [
			sum([g >= 1 for g in [genotypes[s] for s in _samples]])
			for _samples in group_samples]
		group_genotyped = [
			sum([g >= 0 for g in [genotypes[s] for s in _samples]])
			for _samples in group_samples]
		group_cols = ['%d / %d / %d' %
			(group_homozygous[k], group_heterozygous[k], group_genotyped[k])
			for k in range(len(group_homozygous))]

		# Calculate a score for the variant.
		score = float(sum([g > 0 for g in genotypes])) / len(samples)
		score += 1 if info_cols[col_cosmic] != '-' else 0
		score -= 2 if info_cols[col_1000g] != '-' else 0
		
		variant_scores.append(score)
		variant_lines.append('\t'.join(info_cols + group_cols + gt_cols))

		# Keep track of which variants belong to which gene since user wishes
		# to group variants by genes. FIXME: Ignore intergenic variants?
		if rank_method.lower() == 'gene_frequency':
			nearby_genes = cols[col_nearby_genes].split(';')
			for gene in nearby_genes:
				gene_variants[gene].append(len(variant_lines)-1)
			
			
	
	# Print the original headers plus group mutation count headers.
	group_headers = [gname + ' HOM/HET/TOT' for gname in group_names]
	print('\t'.join(headers[:sample_col] + group_headers + samples))

	if rank_method.lower() == 'frequency':
		# Sort the variants based on their individual scores and print them out.
		variant_order = argsort(variant_scores)[::-1]
		for i in variant_order:
			print(variant_lines[i])
	elif rank_method.lower() == 'gene_frequency':
		# First sort the genes based on overall mutation frequency across
		# all amino acids.
		gene_variants = gene_variants.values()
		gene_scores = [sum([max(variant_scores[v], 0) for v in gvariants])
			for gvariants in gene_variants]
		gene_order = argsort(gene_scores)[::-1]
		for k in gene_order:
			gene_variant_scores = [variant_scores[v] for v in gene_variants[k]]
			gene_variant_lines = [variant_lines[v] for v in gene_variants[k]]
			variant_order = argsort(gene_variant_scores)[::-1]
			for i in variant_order:
				print(gene_variant_lines[i])






#######################
# COMMAND LINE PARSER #
#######################

if __name__ == '__main__':
	args = docopt.docopt(__doc__)
	if args['filter']:
		vcf_filter(args['<vcf_file>'], min_qual=int(args['--min-qual']),
			min_alt_control_qual=int(args['--min-alt-control-qual']),
			test_regex=args['--test'], control_regex=args['--control'],
			only_nonsynonymous=args['--nonsynonymous'],
			baseq_p_threshold=float(args['--baseq-alpha']),
			min_alternate=int(args['--min-alt']))
	elif args['summarize']:
		vcf_summarize(args['<vcf_file>'], min_qual=int(args['--min-qual']),
			group_names=args['--group-names'],
			group_regex=args['--group-regex'],
			rank_method=args['--rank'])


